from langgraph.graph import END, START, StateGraph
from langgraph.graph import MessagesState
from langgraph.types import Command
from typing import Annotated, Optional
from typing import Literal
import operator
import asyncio
from pydantic import BaseModel, Field
from langchain_core.messages import MessageLikeRepresentation
from langchain_core.runnables import RunnableConfig
from langchain_core.messages import (
    AIMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
    filter_messages,
    get_buffer_string,
)

from data import AgentState, AgentInputState, SupervisorState, ResearcherState, ResearcherOutputState
from clarify_with_user import clarify_with_user
from write_research_brief import write_research_brief
from supervisor import supervisor
from supervisor_tools import supervisor_tools
from final_report_generation import final_report_generation
from configuration import Configuration
import uuid


###############################################################################
# Supervisor Subgraph Construction
# Creates the supervisor workflow that manages research delegation and coordination
supervisor_builder = StateGraph(
    state_schema=SupervisorState, 
    config_schema=Configuration,
)

# Add supervisor nodes for research management
supervisor_builder.add_node("supervisor", supervisor)           # Main supervisor logic
supervisor_builder.add_node("supervisor_tools", supervisor_tools)  # Tool execution handler

# Define supervisor workflow edges
supervisor_builder.add_edge(START, "supervisor")  # Entry point to supervisor

# Compile supervisor subgraph for use in main workflow
supervisor_subgraph = supervisor_builder.compile()

###############################################################################

# Main Deep Researcher Graph Construction
# Creates the complete deep research workflow from user input to final report
deep_researcher_builder = StateGraph(
    state_schema=AgentState,
    input_schema=AgentInputState, 
    config_schema=Configuration
)

# Add main workflow nodes for the complete research process
deep_researcher_builder.add_node("clarify_with_user", clarify_with_user) # User clarification phase
deep_researcher_builder.add_node("write_research_brief", write_research_brief)     # Research planning phase
deep_researcher_builder.add_node("research_supervisor", supervisor_subgraph)       # Research execution phase
deep_researcher_builder.add_node("final_report_generation", final_report_generation)  # Report generation phase

deep_researcher_builder.add_edge(START, "clarify_with_user")
deep_researcher_builder.add_edge("research_supervisor", "final_report_generation") # Research to report
deep_researcher_builder.add_edge("final_report_generation", END)                   # Final exit point

deep_researcher = deep_researcher_builder.compile()

async def main():
    config = {
        "configurable": {
            "thread_id": str(uuid.uuid4()),
        }
    }
    config["configurable"]["max_researcher_iterations"] = 5

    user_input = "Please summarize weather report for past year?"
    user_input = "Please summarize the weather report for the past year, focusing on temperature trends and precipitation patterns in the United States."
    async for chunk in deep_researcher.astream(
        {"messages": [HumanMessage(content=user_input)]},
        config=config,  # Pass the configuration dictionary here
        stream_mode="updates"
    ):
        pass
        # print("------------------ Graph response --------------------")
        # print(chunk)
        # print("\n")

asyncio.run(main())
